#coding:utf8
'''
Created on 2016年6月12日

@author: zhangjinrui
'''

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid
import matplotlib.animation as animation
from matplotlib.widgets import Button
from DataReader import DataReader
import threading
import time
import matplotlib as mpl
from matplotlib import rc
from scipy.constants.constants import alpha

path = '../../data'
padl = 0.05
padr = 0.1
padb = 0.15
padt = 0.15
pad_each = 0.005
height = 0.8
colorbar_height = 0.03
pad_colorbar = 0.005

class presentInfo:
    
    def __init__(self, f):

        self.figure = f
        self.figure.suptitle(u'集群负载可视化', size = 30, color = 'r')
        self.annos = {}
        self.data_reader = DataReader(path)
        self.host_grid = {}
        self.hosts = self.data_reader.get_hostnames()
        self.init_grids()
        
        
    def init_grids(self):
        self.host_count = self.data_reader.get_host_count()
        self.width = ((1 - padl - padr) - (pad_each * (self.host_count - 1))) / self.host_count
        self.height = height
        for i in range(self.host_count):
            #设计布局
            self.host_grid[self.hosts[i]] = self.add_one_host(i)
            #初始化图标状态
            self.init_one_host_axes(self.hosts[i],i)
        
        #添加一个时间axes
        rect = 0.4 , 0.03, 0.2, 0.06
        self.timeAxes = self.figure.add_axes(rect)
        self.timeAxes.set_xticklabels([0,1.0], visible = False)
        self.timeAxes.set_yticklabels([0,1.0], visible = False)
        self.timeLabel = self.timeAxes.text(0.5,0.5, 'Time: ',va='center', ha='center', size='18')
        
    
            
    def add_one_colorbar(self, axes):
        
        cmap = mpl.colors.ListedColormap([[70.0 / 255, 200.0 / 255, 104.0 /255], 
                                          [1., .8, 0.], 
                                          [1., .4, 0.],
                                          [255.0 / 255, 75.0 / 255, 79.0 /255]])
        cmap.set_under((157.0 / 255, 216.0 / 255, 113.0 / 255))
        cmap.set_over((158.0/ 255, 2.0 / 255, 1.0 / 255))
        
        
        bounds = [0.2, .4, .6, .8, 1.]
        norm = mpl.colors.BoundaryNorm(bounds, cmap.N)
        cb3 = mpl.colorbar.ColorbarBase(axes, cmap=cmap,
                                        norm=norm,
                                        boundaries=[-10] + bounds + [10],
                                        extend='both',
                                        # Make the length of each extension
                                        # the same as the length of the
                                        # interior colors:
                                        extendfrac='auto',
                                        ticks=bounds,
                                        spacing='uniform',
                                        orientation='horizontal')
        axes.set_xticklabels([0,1], visible = False)
        
    
    def add_one_host(self,pos):
        height = (1 - padt - padb) / 4
        
        rect = padl + pos * (self.width + pad_each), padb + 3 * height + colorbar_height + pad_colorbar, self.width, height 
        ax1 = self.figure.add_axes(rect)
        #设置hostName
        ax1.set_title(self.hosts[pos])
        self.fade_xlabel(ax1)
        
        rect = padl + pos * (self.width + pad_each), padb + 2 * height + colorbar_height + pad_colorbar, self.width, height 
        ax2 = self.figure.add_axes(rect)
        self.fade_xlabel(ax2)
        
        rect = padl + pos * (self.width + pad_each), padb + 1 * height + colorbar_height + pad_colorbar, self.width, height 
        ax3 = self.figure.add_axes(rect)
        self.fade_xlabel(ax3)
        
        rect = padl + pos * (self.width + pad_each), padb + 0 * height + colorbar_height + pad_colorbar, self.width, height 
        ax4 = self.figure.add_axes(rect)
        self.fade_xlabel(ax4)
        
        #添加colorbar
        rect = padl + pos * (self.width + pad_each), padb , self.width, colorbar_height
        ax5 = self.figure.add_axes(rect)
        
        #初始化annotition
        self.annos[self.hosts[pos]] = [None,None,None,None,None,None,None]
        
        
        if pos == 0:
            ax1.set_ylabel('CPU')
            ax2.set_ylabel('MEMORY')
            ax3.set_ylabel('DISK')
            ax4.set_ylabel('NET')
            
        if pos != 0 :
            self.fade_ylabel(ax1)
            self.fade_ylabel(ax2)
            self.fade_ylabel(ax3)
            self.fade_ylabel(ax4)
        
        
        return list([ax1,ax2,ax3,ax4,ax5])
    
    def fade_xlabel(self,ax):
        ax.set_xticklabels(range(10), visible = False)
    
    def fade_ylabel(self,ax):
        ax.set_yticklabels(range(10), visible = False)
    
    def add_one_grid(self,pos):
        rect = padl + pos * (self.width + pad_each) , padb, self.width, self.height
        print rect
        grid = AxesGrid(self.figure, rect,  
#         grid = AxesGrid(self.figure, 100 + self.host_count*10 + pos + 1,  
                        nrows_ncols=(4, 1),
                        axes_pad=0.0,
                        share_all=False,
#                         add_all = True
#                         label_mode="L",
#                         cbar_location="top",
#                         cbar_mode="single",
                        )
        return grid
    
    #为了animate..
    def get_data(self):
        while True:
            yield self.data_reader.get_one_set_data()
    
    def update_data(self,data):
        for i in range(len(self.hosts)):
            self.set_data_to_one_host(i, self.hosts[i], data[self.hosts[i]])
#         self.figure.canvas.draw()
#         self.set_data_to_one_host(0, self.hosts[0], data[self.hosts[0]])
        #更新时间
        date = self.get_format_time(data[self.hosts[0]]['timestamp'])
        self.timeLabel.set_text('Time: ' + date)
        
    def get_format_time(self,t):
        t = time.localtime(t)
        return time.strftime('%Y-%m-%d %H:%M:%S',t)
    
    def init_one_host_axes(self, name, pos):
        tx = [0,1,2,3,4,5,6,7,8,9]
        ty = [0,0,0,0,0,0,0,0,0,0]
        
        cpu_axes = self.host_grid[name][0]
        memory_axes = self.host_grid[name][1]
        disk_axes = self.host_grid[name][2]
        net_axes = self.host_grid[name][3]
        color_axes = self.host_grid[name][4]
        
        cpu_axes.plot(tx,ty)
        cpu_axes.get_lines()[0].set_color('r')
        cpu_axes.set_ylim(0,100)
        cpu_axes.set_xlim(0,10)
        cpu_axes.grid(True)
        
        
        memory_axes.plot(tx,ty)
        memory_axes.set_ylim(0,100)
        memory_axes.set_xlim(0,10)
        memory_axes.grid(True)
        
        disk_axes.plot(tx,ty)
        disk_axes.plot(tx,ty)
        disk_axes.set_ylim(-5000,5000)
        disk_axes.set_xlim(0,10)
        disk_axes.grid(True)
        
        net_axes.plot(tx,ty)
        net_axes.plot(tx,ty)
        net_axes.set_ylim(-3000,3000)
        net_axes.set_xlim(0,10)
        net_axes.set_xticks(range(10))
        net_axes.grid(True)
        
        #添加colorbar
        self.add_one_colorbar(color_axes)
        
        #设置legend
        if pos == len(self.hosts) - 1:
#         if pos == 0:
            cpu_axes.legend( [ 'CPU'], loc='center left', bbox_to_anchor=(1, 0.5))
            memory_axes.legend( [ 'MEMORY'], loc='center left', bbox_to_anchor=(1, 0.5))
            disk_axes.legend( [ 'Disk_Read','Disk_Write'], loc='center left', bbox_to_anchor=(1, 0.5))
            net_axes.legend( [ 'Net_Send','Net_Recieve'], loc='center left', bbox_to_anchor=(1, 0.5))
            
        
    def set_data_to_one_host(self, pos, name, dat):
        
        cpu_axes = self.host_grid[name][0]
        memory_axes = self.host_grid[name][1]
        disk_axes = self.host_grid[name][2]
        net_axes = self.host_grid[name][3]
        colorbar_axes = self.host_grid[name][4]
        
        #更新CPU的值
        line, = cpu_axes.get_lines()
        y = list(line.get_ydata())
        y.append(dat['cpu_all_percent'])
        if len(y) > 10:
            del y[0]
        x = range(len(y))
        line.set_data(x , y)
        #更新annotation
        maxy = max(y)
        maxx = y.index(maxy)
        xtext = maxx + 2
        ytext = maxy + 20 if maxy < 50 else maxy - 20
        an = self.add_annotation(cpu_axes, (maxx, maxy), (xtext, ytext), str(maxy) + '%')
        if self.annos[self.hosts[pos]][0] is not None:
            (self.annos[self.hosts[pos]][0]).remove()
        self.annos[self.hosts[pos]][0] = an
        
        #更新内存的值
        
        line, = memory_axes.get_lines()
        y = list(line.get_ydata())
        x = range(len(y))
        y.append(dat['memory_percent'])
        if len(y) > 10:
            del y[0]
        
#         line.set_data(x, y)
        #这个地方写的有点奇怪
        memory_axes.clear()
        memory_axes.set_ylim(0,100)
        memory_axes.set_xlim(0,10)
        memory_axes.grid(True)
        memory_axes.plot(x,y)
        self.fade_xlabel(memory_axes)
        if pos != 0:
            self.fade_ylabel(memory_axes)
        if pos == len(self.hosts) - 1:
            memory_axes.legend( [ 'MEMORY'], loc='center left', bbox_to_anchor=(1, 0.5))
            
        #更新annotation
        maxy = max(y)
        maxx = y.index(maxy)
        xtext = maxx + 2
        ytext = maxy + 20 if maxy < 50 else maxy - 20
        an = self.add_annotation(memory_axes, (maxx, maxy), (xtext, ytext), str(maxy) + '%')
        if self.annos[self.hosts[pos]][1] is not None:
            (self.annos[self.hosts[pos]][1]).remove()
        self.annos[self.hosts[pos]][1] = an
        #更新fill
#         for coll in (memory_axes.collections):
#             memory_axes.collections.remove(coll)
        
        
        memory_axes.fill_between(x, 0, y, alpha = 0.3, color = 'r')
        
        
        #更新磁盘的值
        lines = disk_axes.get_lines()
        #先更新读入速度
        y = list(lines[0].get_ydata())
        y.append(dat['disk_read_speed'])
        if len(y) > 10:
            del y[0]
        x = range(len(y))
        lines[0].set_data(x , y)
        #更新Text
        maxy = max(y)
        if self.annos[self.hosts[pos]][2] is None:
            self.annos[self.hosts[pos]][2] = disk_axes.text(0.2,4000, 'R: ' + str(maxy) + 'Kb/s')
        else:
            self.annos[self.hosts[pos]][2].set_text('R:' + str(maxy) + 'Kb/s')
        
        #更新写入速度
        y = list(lines[1].get_ydata())
        y.append(- dat['disk_write_speed'])
        if len(y) > 10:
            del y[0]
        x = range(len(y))
        lines[1].set_data(x , y)
        #更新Text
        maxy = min(y)
        if self.annos[self.hosts[pos]][3] is None:
            self.annos[self.hosts[pos]][3] = disk_axes.text(0.2,-4000, 'W: ' + str(-maxy) + 'Kb/s')
        else:
            self.annos[self.hosts[pos]][3].set_text('W:' + str(-maxy) + 'Kb/s')

        
        #更新网络的值
        lines = net_axes.get_lines()
        #更新发送速度
        y = list(lines[0].get_ydata())
        y.append(dat['net_send_speed'])
        if len(y) > 10:
            del y[0]
        x = range(len(y))
        lines[0].set_data(x , y)
        #更新Text
        maxy = max(y)
        if self.annos[self.hosts[pos]][4] is None:
            self.annos[self.hosts[pos]][4] = net_axes.text(0.2,2500, 'S: ' + '%.2f'%maxy + 'Kb/s')
        else:
            self.annos[self.hosts[pos]][4].set_text('S:' + '%.2f'%maxy + 'Kb/s')
        
        #更新接受速度
        y = list(lines[1].get_ydata())
        y.append(- dat['net_rec_speed'])
        if len(y) > 10:
            del y[0]
        x = range(len(y))
        lines[1].set_data(x , y)
        #更新Text
        maxy = -min(y)
        if self.annos[self.hosts[pos]][5] is None:
            self.annos[self.hosts[pos]][5] = net_axes.text(0.2,-2500, 'R: ' + '%.2f'%maxy + 'Kb/s')
        else:
            self.annos[self.hosts[pos]][5].set_text('S:' + '%.2f'%maxy + 'Kb/s')
        
        
        #更新colorbar
        xv = self.cal_pressure(dat)
        an = self.add_colorbar_annotation(colorbar_axes, (xv, 1), (xv, -1), str(xv))
        if self.annos[self.hosts[pos]][6] is not None:
            (self.annos[self.hosts[pos]][6]).remove()
        self.annos[self.hosts[pos]][1] = an
        
    def cal_pressure(self,dat):
        return dat['cpu_all_percent'] * 1.0 /100
        
    
    def timer_update_data(self):
        print time.time()
        data = self.data_reader.get_one_set_data()
        for i in range(len(self.hosts)):
            self.set_data_to_one_host(i, self.hosts[i], data[self.hosts[i]])
        self.figure.canvas.draw()
        
        self.timer1 = threading.Timer(1, self.timer_update_data)
        self.timer1.start()
    
    def but_on_click(self, event):
        self.an.remove()
    
    def add_colorbar_annotation(self, ax, xy1, xytext1, text):
        an = ax.annotate(text, xy=xy1, xycoords='data',
                xytext=xytext1,
                arrowprops=dict(arrowstyle="-",),
#                 verticalalignment='top',
                horizontalalignment='center'
                )
        return an
    
    def add_annotation(self, ax, xy1, xytext1 ,text):
        an = ax.annotate(text, xy=xy1, xycoords='data',
                xytext=xytext1,
                arrowprops=dict(arrowstyle="->",
                                connectionstyle="angle3,angleA=0,angleB=-90"),
                )
        return an
                    
        
if __name__ == '__main__':
    
    f = plt.figure(1, figsize=(20,10))
    p = presentInfo(f)
    
    
    ani = animation.FuncAnimation(f, p.update_data, p.get_data, interval=1000)
        
#     axnext = plt.axes([0.45,0.1,0.05,0.05])
#     but = Button(axnext, 'Next')
#     but.on_clicked(p.but_on_click)
    
    plt.show()